import os

import torch
from matplotlib import pyplot as plt
from matplotlib.pyplot import figure
from torch.utils.data import DataLoader, Dataset
import numpy as np
import matplotlib.dates as mdates
import datetime as dt
import time


class MyDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __getitem__(self, item):
        return self.data[item]

    def __len__(self):
        return len(self.data)


def write_log(log, txt_list):
    print(log)
    txt_list.append(log + '\r\n')


def process(data, batch_size, shuffle, interval, pred_size, n_output):
    load = data
    lab = data[:, -2:]
    load = load.tolist()
    lab = lab.tolist()
    seq = []
    for i in range(len(data) - interval - pred_size):
        train_seq = []
        train_label = []
        for j in range(i, i + interval):
            x = load[j]
            train_seq.append(x)
        for j in range(i + interval, i + interval + pred_size):
            x = lab[j][n_output]
            train_label.append(x)
        train_seq = torch.FloatTensor(train_seq)
        # train_label = torch.FloatTensor(train_label).view(-1)

        train_label = torch.FloatTensor(train_label)
        # print(train_seq.shape, train_label.shape)

        seq.append((train_seq, train_label))

    seq = MyDataset(seq)
    seq = DataLoader(dataset=seq, batch_size=batch_size, shuffle=shuffle, num_workers=0, drop_last=True)
    return seq


def plot(fold_name, title, t_pre, y_pre, t_real, y_real):
    figure(figsize=(12.8, 9.6))
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    plt.plot(t_pre, y_pre, color='blue', label='预测曲线')
    plt.plot(t_real, y_real, color='red', label='实际曲线')
    plt.title(f'{title}', fontsize=20)
    plt.setp(plt.gca().xaxis.get_majorticklabels(), 'rotation', 30, 'fontsize', 10)
    plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%m-%d %H:%M:%S'))
    plt.xticks(fontsize=15)
    plt.yticks(fontsize=20)
    plt.legend(fontsize=20)
    plt.savefig(f'{fold_name}/{title}.png')


def t_plot(fold_name, title, lstm_model, test_num, t_all, Dte, data, device, pred_size, interval, txt_list):
    os.makedirs(f'{fold_name}/test', exist_ok=True)
    lstm_model.eval()
    ts = test_num + interval
    te = test_num + interval + pred_size + len(Dte) - 1
    t_test = t_all[ts:te]
    y_t = data[ts:te]
    y_pred_list = np.zeros((len(Dte), len(Dte) + pred_size - 1), dtype=float)
    count = 0
    ts = time.perf_counter()
    for seq, _ in Dte:
        seq = seq.to(device)
        with torch.no_grad():
            y_pred = lstm_model(seq)
            y_pred = y_pred.cpu().detach().numpy()

            y_pred_list[count, count:count + len(y_pred[0])] = y_pred[0]

            # y_pred_list += y_pred.tolist()
        count += 1
    y_re = []
    for i in range(len(Dte) + pred_size - 1):
        col = y_pred_list[:, i]
        r_col = col[np.where(col != 0)]
        y_re.append(sum(r_col) / len(r_col))
    te = time.perf_counter()
    wt = str(round(1000 * (te - ts) / len(Dte), 3))
    write_log('time per prediction: ' + wt + ' ms', txt_list)
    y_re = np.array(y_re)
    plot(f'{fold_name}/test', f'{title}-过程线与预测线', t_test, y_re, t_all, data)
    plot(f'{fold_name}/test', f'{title}-预测部分', t_test, y_re, t_test, y_t)
    return y_re, y_t
